#!/usr/bin/env python3
import numpy as np
import base64
import io
import json

def encode_array(arr):
    buffer = io.BytesIO()
    np.save(buffer, arr, allow_pickle=False)
    buffer.seek(0)
    return base64.b64encode(buffer.read()).decode("utf-8")

def main():
    # 固定路径
    npy_file = "../out/layer_features_epoch_01/image_0000/conv1_input.npy"
    out_json = "save_tensors_inputs_list.json"

    # 读取 npy 文件
    arr = np.load(npy_file)

    # 构造 Polygraphy 可识别结构
    sample = {
        "input": {
            "array": encode_array(arr),
            "polygraphy_class": "ndarray",
        }
    }

    # 写入 JSON 数组格式
    with open(out_json, "w") as f:
        json.dump([sample], f)

    print(f"✅ 已保存 Polygraphy 输入文件: {out_json}")
    print(f"📐 shape = {arr.shape}, dtype = {arr.dtype}")

if __name__ == "__main__":
    main()
